import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO

import supervision as sv

model = YOLO("yolon11-seg.pt")
image = cv2.imread(<your_image_path>)


def callback(image_slice: np.ndarray) -> sv.Detections:
    results = model(image_slice)[0]
    return sv.Detections.from_ultralytics(results)

# slicer = sv.InferenceSlicer(callback=callback, slice_wh=(960, 960), overlap_ratio_wh=None, overlap_wh=(50, 50), overlap_filter="non_max_merge", iou_threshold=0.1, match_metric="IOU")
slicer = sv.InferenceSlicer(callback=callback, slice_wh=(960, 960), overlap_ratio_wh=None, overlap_wh=(50, 50), overlap_filter="non_max_merge", iou_threshold=0.1, match_metric="IOS")
# slicer = sv.InferenceSlicer(callback=callback, slice_wh=(960, 960), overlap_ratio_wh=None, overlap_wh=(50, 50), overlap_filter="non_max_suppression", iou_threshold=0.1, match_metric="IOU")
# slicer = sv.InferenceSlicer(callback=callback, slice_wh=(960, 960), overlap_ratio_wh=None, overlap_wh=(50, 50), overlap_filter="non_max_suppression", iou_threshold=0.1, match_metric="IOS")
detections = slicer(image)

box_annotator = sv.BoxAnnotator()
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [f"{class_name} {confidence:.2f}" for class_name, confidence in zip(detections["class_name"], detections.confidence)]

annotated_image = box_annotator.annotate(scene=image, detections=detections)
annotated_image = mask_annotator.annotate(scene=annotated_image, detections=detections)
annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
Image.fromarray(annotated_image[..., ::-1]).save("test.jpg")